import numpy as np
if __name__ == "__main__":
    import misc
else:
    from . import misc

class DecisionStump:
    def __init__(self):
        self.splited = []
        self.col = -1
        

    def train(self, xs, cs, ws):
        cs = np.array(cs)

        minerr = np.inf

        for col in range(len(xs[0,:])):
            indices = xs[:,col].argsort()
            
            splits = []
            s = [xs[indices[0], col], indices[0]]
            start = 0

            for i in range(1,len(indices)):
                if cs[indices[i]] != s[1] and s[0] != xs[indices[i],col]:
                    s[1] = misc.sort_dict(misc.frequency(cs[indices[start:i]]))[-1][0]

                    splits.append(s)
                    s = [xs[indices[i],col], cs[indices[i]]]
                    start = i

                else:
                    s[0] = xs[indices[i],col]
                
            s[0] = np.inf
            s[1] = misc.sort_dict(misc.frequency(cs[indices[start:]]))[-1][0]
            splits.append(s)

            test_cs = np.array([self._classify_single(col, splits, x) for x in xs])
                        
            test_result = np.ones(len(cs))
            test_result[test_cs == np.array(cs)] = 0
                        
            err = test_result.dot(ws)
            if err < minerr:
                minerr = err
                self.splited = splits
                self.col = col
        #print(self.splited)
        #print(self.col)
           
    def classify(self, xs):
        return [self.classify_single(x) for x in xs]

    def classify_single(self, x):
        return self._classify_single(self.col, self.splited, x)

    def _classify_single(self, col, splited, x):

        for s in splited:
            if x[col] <= s[0]:
                return s[1]
        return None

    
        


class DecisionStump1:
    def __init__(self):
        self.stepsize = 10
        self.split_col = -1
        self.thresh = 0
        self.c_choiced = [None, None]


    def train(self, xs, c, D):
        cset = set(c)
        #print(D)
        
        min_err = np.inf
        for col in range(len(xs[0,:])):
            (min, max) = (np.min(xs[:,col]), np.max(xs[:,col]))
            step = (max - min) / float(self.stepsize)


            for i in range(-1, self.stepsize+1):
                thresh = min + i*step
                

                for c1 in cset:
                    _c = np.array(c)
                    _c[xs[:,col] <= thresh] = c1
                    for c2 in cset.difference(set([c1])):
                        _c[xs[:,col] > thresh] = c2
                        
                        test = np.ones(len(c))
                        test[_c == np.array(c)] = 0
                        
                        err = test.dot(D)
         
                        if err < min_err:
                            min_err = err
                            self.split_col = col
                            self.thresh = thresh
                            self.c_choiced = [c1, c2]
        
                        


    def classify(self, xs):

        return [self.c_choiced[0] if x[self.split_col] <= self.thresh else self.c_choiced[1] for x in xs]



class AdaBoost:
    def __init__(self, weaker_type, size):
        self.size = size
        self.weaker_type = weaker_type
        self.weakers = []

    def train(self, xs, cs, params):
        (n,m) = np.shape(xs)

        D = np.ones(n) / n

        for i in range(self.size):
            weaker = self.weaker_type()
            weaker.train(xs, cs, D)


            test_cs = weaker.classify(xs)
            
            correct_cs = np.ones(n)
            correct_cs[np.array(cs) != np.array(test_cs)] = 0

            failed_cs = np.ones(n)
            failed_cs[np.array(cs) == np.array(test_cs)] = 0

            if np.sum(failed_cs*D) > 0:
                alpha = 1/2 * np.log(np.sum(correct_cs*D) / np.sum(failed_cs*D))

                D = D*np.exp(alpha*(failed_cs-correct_cs))
                D = D/np.sum(D)

                self.weakers.append([weaker, alpha])
            else:
                self.weakers.append([weaker, 1])
                break
       
       
        print(self.weakers)

    def classify(self, xs):
        return [self.classify_single(x) for x in xs ]

    def classify_single(self, x):
        ws_c = {}
        for weaker in self.weakers:
            [c] = weaker[0].classify([x])
            ws_c[c] = ws_c.get(c, 0) + weaker[1]
        return misc.sort_dict(ws_c)[-1][0]

def _test_decision_stump():
    xs = np.array([[1,2,3],[4,5,6],[7,8,9],[9,8,7]])
    cs = np.array(['a','b','c','c'])

    stump = DecisionStump()
    stump.train(xs, cs, np.array([1,1,1,1]).T)
    print(stump.classify(np.array([[2,2,3],[8,8,8]])))


def _test_adaboost():
    
    xs = np.array([[1,2,3],[4,5,6],[7,8,9],[9,8,7]])
    cs = np.array(['a','b','c','c'])

    ab = AdaBoost(DecisionStump, 2)

    ab.train(xs, cs, {})
    print("--------------")
    print(ab.classify(np.array([[2,2,3],[8,8,8]])))

if __name__ == "__main__":
    _test_adaboost()

